import socket
import struct
import ipaddress

def ipv4_checksum(header_bytes: bytes) -> int:
    """
    计算IPv4头部的16位校验和（one's complement）。
    输入长度应为偶数（若为奇数会补0）。
    返回值是一个 0..0xffff 的整数。
    """
    if len(header_bytes) % 2 == 1:
        header_bytes += b'\x00'
    s = 0
    # 按16位累加
    for i in range(0, len(header_bytes), 2):
        word = (header_bytes[i] << 8) + header_bytes[i+1]
        s += word
    # 处理进位，高于16位的部分加回低16位，直到没有进位
    while s >> 16:
        s = (s & 0xFFFF) + (s >> 16)
    # 取反得到校验和
    return (~s) & 0xFFFF


def build_ipv4_header(src_ip: str,
                      dst_ip: str,
                      payload_len: int,
                      protocol: int = 17,
                      identification: int = 1,
                      ttl: int = 64,
                      tos: int = 0,
                      flags: int = 0,
                      frag_offset: int = 0,
                      ihl: int = 5) -> bytes:
    """
    构造一个标准 IPv4 头（IHL=5，20 字节），并计算头部校验和。
    参数:
      - src_ip, dst_ip: IPv4 字符串，如 "192.168.2.10"
      - payload_len: 上层负载长度（字节），函数会自动计算 total_length = ihl*4 + payload_len
      - protocol: 上层协议号（例如 UDP=17, TCP=6, ICMP=1）
      - identification: 16-bit 标识
      - ttl: 生存时间（0-255）
      - tos: Type of Service / DSCP
      - flags: 3-bit (0..7)，例如 DF=2, MF=1
      - frag_offset: 13-bit 片偏移
      - ihl: IPv4 头长（>=5），本函数默认5（即无选项）
    返回:
      - 20 字节的 IPv4 头 (bytes)
    会对输入做基本校验（IP 格式、字段范围等），不支持 IHL>5 的选项扩展。
    """
    # 参数校验
    try:
        ipaddress.IPv4Address(src_ip)
        ipaddress.IPv4Address(dst_ip)
    except Exception as e:
        raise ValueError("src_ip 或 dst_ip 不是合法的 IPv4 地址") from e

    if not (5 <= ihl <= 15):
        raise ValueError("IHL 必须在 5..15 之间（通常为5）")
    if not (0 <= flags <= 0x7):
        raise ValueError("flags 必须是 0..7（3 位）")
    if not (0 <= frag_offset <= 0x1FFF):
        raise ValueError("frag_offset 必须是 0..8191（13 位）")
    if not (0 <= tos <= 0xFF and 0 <= ttl <= 0xFF and 0 <= protocol <= 0xFF):
        raise ValueError("tos, ttl, protocol 必须在 0..255")

    ver = 4
    ver_ihl = (ver << 4) + ihl                # 1 字节: 版本(4bit) + IHL(4bit)
    total_length = ihl*4 + int(payload_len)  # 16-bit 总长度（头+数据）

    flags_frag = (flags << 13) | (frag_offset & 0x1FFF)  # 16-bit Flags(3) + Fragment offset(13)

    # 先把 checksum 字段填 0，再计算校验和
    src_packed = socket.inet_aton(src_ip)
    dst_packed = socket.inet_aton(dst_ip)

    # 按网络字节序打包（checksum先填0）
    header_wo_checksum = struct.pack('!BBHHHBBH4s4s',
                                     ver_ihl,
                                     tos,
                                     total_length,
                                     identification & 0xFFFF,
                                     flags_frag & 0xFFFF,
                                     ttl & 0xFF,
                                     protocol & 0xFF,
                                     0,                 # checksum 占位
                                     src_packed,
                                     dst_packed)

    chksum = ipv4_checksum(header_wo_checksum)

    header = struct.pack('!BBHHHBBH4s4s',
                         ver_ihl,
                         tos,
                         total_length,
                         identification & 0xFFFF,
                         flags_frag & 0xFFFF,
                         ttl & 0xFF,
                         protocol & 0xFF,
                         chksum,
                         src_packed,
                         dst_packed)
    return header


if __name__ == "__main__":
    src = "5.23.1.10"
    dst = "5.23.2.10"
    payload_len = 82 - 20
    hdr = build_ipv4_header(src, dst, payload_len, protocol=0xfd, identification=0x3acc, ttl=0x80)
    print("header (bytes):", hdr)
    print("header (hex):", hdr.hex())
    # 可插入以太帧后用 raw socket 发送
